import sys
import csv
import time
import socket
import logging
from datetime import datetime
from ping3 import ping
from PyQt5.QtWidgets import (QApplication, QMainWindow, QWidget, QTableView, QHeaderView,
                            QLabel, QLineEdit, QPushButton, QComboBox, QFileDialog,
                            QVBoxLayout, QHBoxLayout, QMessageBox, QStyleFactory)
from PyQt5.QtCore import (Qt, QAbstractTableModel, QModelIndex, QThreadPool,
                         pyqtSignal, QTimer, QRunnable, QObject)
from PyQt5.QtGui import QColor, QIntValidator
from logging.handlers import RotatingFileHandler

# 配置日志系统
def setup_logging():
    logger = logging.getLogger()
    logger.setLevel(logging.DEBUG)

    # 文件日志（自动轮换，最大5个文件，每个10MB）
    file_handler = RotatingFileHandler(
        "network_monitor.log",
        maxBytes=10*1024*1024,
        backupCount=5,
        encoding="utf-8"
    )
    file_formatter = logging.Formatter(
        "%(asctime)s - %(name)s - %(levelname)s - %(threadName)s - %(message)s"
    )
    file_handler.setFormatter(file_formatter)

    # 控制台日志
    console_handler = logging.StreamHandler()
    console_formatter = logging.Formatter("%(asctime)s - %(levelname)s - %(message)s")
    console_handler.setFormatter(console_formatter)

    logger.addHandler(file_handler)
    logger.addHandler(console_handler)

    # 捕获未处理异常
    def handle_exception(exc_type, exc_value, exc_traceback):
        if issubclass(exc_type, KeyboardInterrupt):
            sys.__excepthook__(exc_type, exc_value, exc_traceback)
            return
        logger.critical("未捕获异常", exc_info=(exc_type, exc_value, exc_traceback))

    sys.excepthook = handle_exception

setup_logging()
logging = logging.getLogger(__name__)

class ResultModel(QAbstractTableModel):
    COLUMNS = ["目标", "类型", "端口", "状态", "延迟(ms)", "成功", "总数", "最后失败"]
    STATUS_COLOR_MAP = {
        "成功": QColor("#2ecc71"),
        "失败": QColor("#e74c3c"),
        "待检测": QColor("#95a5a6")
    }
    DELAY_COLOR_MAP = {
        "low": QColor("#2ecc71"),    # 0-100ms
        "medium": QColor("#f1c40f"), # 100-300ms
        "high": QColor("#e74c3c")    # 300ms+
    }
    DELAY_LOW = 30
    DELAY_MEDIUM = 60

    def __init__(self):
        super().__init__()
        self._data = []
        self._target_map = {}  # { (target, type, port): row_index }
        self._raw_delay = []  # 保存原始延迟数值
    def rowCount(self, parent=QModelIndex()):
        return len(self._data)

    def columnCount(self, parent=QModelIndex()):
        return len(self.COLUMNS)

    def data(self, index, role=Qt.DisplayRole):
        if not index.isValid():
            return None

        row = index.row()
        col = index.column()
        item = self._data[row]
        if role == Qt.DisplayRole:
            return item[col]
        elif role == Qt.ForegroundRole and col == 3:  # 状态列颜色
            return self.STATUS_COLOR_MAP.get(item[col], QColor(Qt.black))
        # 延迟列颜色
        elif role == Qt.ForegroundRole and col == 4:
            try:
                delay = float(self._raw_delay[row])
                if delay <= self.DELAY_LOW:
                    return self.DELAY_COLOR_MAP["low"]
                elif self.DELAY_LOW < delay <= self.DELAY_MEDIUM:
                    return self.DELAY_COLOR_MAP["medium"]
                else:
                    return self.DELAY_COLOR_MAP["high"]
            except Exception as e:
                logging.error(f"延迟值错误: {e}")
                return QColor(Qt.gray)
        elif role == Qt.TextAlignmentRole:
            return Qt.AlignCenter
        return None

    def headerData(self, section, orientation, role=Qt.DisplayRole):
        if orientation == Qt.Horizontal and role == Qt.DisplayRole:
            return self.COLUMNS[section]
        return None

    def add_target(self, target_info):
        key = (target_info["target"], target_info["type"], target_info["port"])
        if key in self._target_map:
            return False
        logging.debug(f"self._target_map: {self._target_map}")

        row = self.rowCount()
        self.beginInsertRows(QModelIndex(), row, row)
        self._data.append([
            target_info["target"],
            target_info["type"],
            str(target_info["port"]) if target_info["port"] != 0 else "-",
            "待检测",  # 状态
            "0",      # 延迟
            0,        # 成功次数
            0,        # 总次数
            "无"       # 最后失败时间
        ])
        self._target_map[key] = row
        self._raw_delay.append(0)  # 初始化延迟数据
        self.endInsertRows()
        return True

    def update_result(self, key, success, delay):
        row = self._target_map.get(key)
        if row is None:
            return


        # 更新数据
        self._data[row][3] = "成功" if success else "失败"
        self._data[row][4] = f"{delay:.2f}" if success else "0"
        self._data[row][5] += 1 if success else 0
        self._data[row][6] += 1
        if not success:
            self._data[row][7] = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
        # 保存原始延迟数值
        self._raw_delay[row] = f"{delay:.2f}" if success else "0"
        # 通知视图更新
        self.dataChanged.emit(
            self.index(row, 4), 
            self.index(row, 7),
            [Qt.DisplayRole, Qt.ForegroundRole]
        )

    def target_exists(self, target_info):
        key = (target_info["target"], 
              target_info["type"], 
              target_info["port"])
        return key in self._target_map
    
    def remove_rows(self, rows):
        # 按逆序删除行以避免索引变化问题
        for row in sorted(rows, reverse=True):
            if 0 <= row < len(self._data):
                # 获取被删除的目标信息
                target = self._data[row][0]
                type_ = self._data[row][1]
                port = 0 if self._data[row][2] == "-" else int(self._data[row][2])
                key = (target, type_, port)
                logging.debug(f"删除目标: {key}")
                
                # 更新目标映射
                del self._target_map[key]
                # 调整其他行的索引
                for k in self._target_map:
                    if self._target_map[k] > row:
                        self._target_map[k] -= 1

                # 删除数据
                self.beginRemoveRows(QModelIndex(), row, row)
                del self._data[row]
                del self._raw_delay[row]
                self.endRemoveRows()

class DetectionSignals(QObject):
    result_ready = pyqtSignal(tuple, bool, float)  # (key, success, delay)

class DetectionTask(QRunnable):
    def __init__(self, target_info):
        super().__init__()
        self.target_info = target_info
        self.signals = DetectionSignals()
        self._active = True
        self.setAutoDelete(False)  # 保持任务对象存活
        self.interval = target_info["interval"] 

    def run(self):
        while self._active:
            try:
                key = (self.target_info["target"], 
                      self.target_info["type"], 
                      self.target_info["port"])
                
                if self.target_info["type"] == "PING":
                    success, delay = self.ping_test()
                elif self.target_info["type"] == "TCP":
                    success, delay = self.tcp_test()
                elif self.target_info["type"] == "UDP":
                    success, delay = self.udp_test()
                else:
                    continue

                self.signals.result_ready.emit(key, success, delay)

                time.sleep(int(self.interval))  # 每个目标检测间隔
            except Exception as e:
                logging.error(f"检测异常: {str(e)}")

    def stop(self):
        self._active = False

    def ping_test(self):
        try:
            delay = ping(self.target_info["target"], timeout=1, unit='ms')
            logging.debug(f"PING {self.target_info['target']} 延迟: {delay} ms")
            return (True, delay) if delay else (False, 0)
        except Exception as e:
            return False, 0

    def tcp_test(self):
        try:
            start = time.time()
            with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
                s.settimeout(1)
                s.connect((self.target_info["target"], self.target_info["port"]))
            logging.debug(f"TCP {self.target_info['target']}:{self.target_info['port']} 连接成功")
            return True, (time.time() - start) * 1000
        except Exception as e:
            return False, 0

    def udp_test(self):
        try:
            start = time.time()
            with socket.socket(socket.AF_INET, socket.SOCK_DGRAM) as s:
                s.settimeout(1)
                s.sendto(b"", (self.target_info["target"], self.target_info["port"]))
                s.recvfrom(1024)
            return True, (time.time() - start) * 1000
        except:
            return False, 0

class MainWindow(QMainWindow):
    def __init__(self):
        super().__init__()
        self.setWindowTitle("网络监控工具")
        self.setMinimumSize(800, 600)
        self.thread_pool = QThreadPool.globalInstance()
        self.active_tasks = {}  # {key: DetectionTask}
        self._init_ui()
        self.csv_writer = CSVWriter()

    def _init_ui(self):
        main_widget = QWidget()
        self.setCentralWidget(main_widget)
        layout = QVBoxLayout(main_widget)

        # 输入区域
        input_layout = QHBoxLayout()
        self.target_input = QLineEdit()
        self.target_input.setPlaceholderText("输入目标地址")
        self.target_input.setMinimumWidth(200)  # 最小宽度保障可读性
        self.type_combo = QComboBox()
        self.type_combo.addItems(["PING", "TCP", "UDP"])
        self.type_combo.currentTextChanged.connect(self.update_port_input_state)
        self.port_input = QLineEdit()
        self.port_input.setPlaceholderText("端口")
        self.port_input.setValidator(QIntValidator(1, 65535, self))
        self.port_input.setFixedWidth(80)  # 刚好显示最大端口号65535
        self.add_btn = QPushButton("添加目标")
        self.add_btn.clicked.connect(self._add_target)
        self.type_combo.setFixedWidth(100)
        self.add_btn.setFixedWidth(100)

        input_layout.addWidget(QLabel("目标:"))
        input_layout.addWidget(self.target_input)
        input_layout.addWidget(QLabel("类型:"))
        input_layout.addWidget(self.type_combo)
        input_layout.addWidget(QLabel("端口:"))
        input_layout.addWidget(self.port_input)
        input_layout.addWidget(self.add_btn)

        self.interval_input = QLineEdit()
        self.interval_input.setText("1")
        self.interval_input.setFixedWidth(60)
        self.interval_input.setValidator(QIntValidator(1, 3600, self))
        input_layout.addWidget(QLabel("间隔(s):"))
        input_layout.addWidget(self.interval_input)

        # 控制区域
        control_layout = QHBoxLayout()
        self.start_btn = QPushButton("开始检测")
        self.start_btn.clicked.connect(self.start_detection)
        self.stop_btn = QPushButton("停止检测")
        self.stop_btn.clicked.connect(self.stop_detection)
        self.stop_btn.setEnabled(False)
        self.import_btn = QPushButton("批量导入")
        self.import_btn.clicked.connect(self.import_targets)
        # 在控制布局中添加删除按钮
        self.delete_btn = QPushButton("删除选中")
        self.delete_btn.clicked.connect(self.delete_selected)
        control_layout.addWidget(self.delete_btn)

        control_layout.addWidget(self.start_btn)
        control_layout.addWidget(self.stop_btn)
        control_layout.addWidget(self.import_btn)

        # 结果表格
        self.table_view = QTableView()
        self.model = ResultModel()
        self.table_view.setModel(self.model)
        # self.table_view.horizontalHeader().setSectionResizeMode(QHeaderView.Stretch)
        self.table_view.setSelectionBehavior(QTableView.SelectRows)

        # 配置列宽
        header = self.table_view.horizontalHeader()
        header.setDefaultAlignment(Qt.AlignCenter)
        header.setSectionsMovable(True)  # 允许拖动列顺序

        
        # 设置各列调整策略和初始宽度
        width_config = [
            (0, QHeaderView.Interactive, 250),   # 目标地址
            (1, QHeaderView.Fixed, 80),          # 类型
            (2, QHeaderView.Fixed, 70),          # 端口
            (3, QHeaderView.ResizeToContents, 100), # 状态
            (4, QHeaderView.Interactive, 120),   # 延迟
            (5, QHeaderView.Fixed, 80),          # 成功次数
            (6, QHeaderView.Fixed, 80),          # 总次数
            (7, QHeaderView.Stretch, 200)        # 最后失败时间
        ]

        for col, mode, width in width_config:
            header.setSectionResizeMode(col, mode)
            if mode != QHeaderView.Stretch:
                header.resizeSection(col, width)

                # 优化显示
        self.table_view.setAlternatingRowColors(True)
        self.table_view.setSortingEnabled(True)
        self.table_view.setShowGrid(False)
        self.table_view.setWordWrap(False)
        
        # 在控制布局中添加列宽重置按钮（新增）
        self.reset_col_btn = QPushButton("重置列宽")
        self.reset_col_btn.clicked.connect(self.reset_column_width)
        control_layout.addWidget(self.reset_col_btn)

        # 布局组装
        layout.addLayout(input_layout)
        layout.addLayout(control_layout)
        layout.addWidget(self.table_view)

    def import_targets(self):
        path, _ = QFileDialog.getOpenFileName(
            self, "选择CSV文件", "", "CSV文件 (*.csv)")
        
        if not path:
            return

        success_count = 0
        error_count = 0
        error_log = []
        
        try:
            with open(path, 'r', encoding='utf-8') as f:
                reader = csv.DictReader(f)
                
                # 验证必须字段
                required_fields = {'target', 'type', 'port'}
                if missing := required_fields - set(reader.fieldnames):
                    logging.error(f"缺少必要列: {', '.join(missing)}")
                    raise ValueError(f"缺少必要列: {', '.join(missing)}")

                for line_num, row in enumerate(reader, 2):  # 从第2行开始计数
                    try:
                        # 数据清洗
                        target = row['target'].strip()
                        type_ = row['type'].strip().upper()
                        port_str = row['port'].strip()

                        # 基础验证
                        if not target:
                            raise ValueError("目标地址不能为空")
                            
                        if type_ not in ('PING', 'TCP', 'UDP'):
                            raise ValueError("检测类型必须为 PING/TCP/UDP")

                        # 处理端口
                        port = 0
                        if type_ != 'PING':
                            if not port_str:
                                raise ValueError("TCP/UDP检测需要端口号")
                                
                            # 处理可能的浮点格式 (如 80.0)
                            if '.' in port_str:
                                port_str = port_str.split('.')[0]
                                
                            if not port_str.isdigit():
                                raise ValueError("端口必须为整数")
                                
                            port = int(port_str)
                            if not (1 <= port <= 65535):
                                raise ValueError("端口范围应为1-65535")

                        # 添加目标
                        if self.model.add_target({
                            "target": target,
                            "type": type_,
                            "port": port
                        }):
                            success_count += 1
                        else:
                            raise ValueError("目标已存在")

                    except Exception as e:
                        error_count += 1
                        error_log.append(
                            f"行号 {line_num}: {str(e)} | 数据: {dict(row)}"
                        )
                        logging.error(f"CSV导入错误: {error_log[-1]}")

            # 显示导入结果
            msg = QMessageBox()
            msg.setIcon(QMessageBox.Information)
            msg.setText(
                f"成功导入: {success_count}条\n"
                f"失败: {error_count}条"
            )
            
            if error_log:
                error_path = f"import_errors_{datetime.now():%Y%m%d_%H%M%S}.log"
                with open(error_path, 'w', encoding='utf-8') as f:
                    f.write("\n".join(error_log))
                msg.setDetailedText(f"错误日志已保存至: {error_path}")
                
            msg.setWindowTitle("导入结果")
            msg.exec_()

        except UnicodeDecodeError:
            QMessageBox.critical(self, "编码错误", "请使用UTF-8编码的文件")
        except Exception as e:
            QMessageBox.critical(self, "导入失败", f"文件读取失败: {str(e)}")

    def update_port_input_state(self, type_):
        self.port_input.setEnabled(type_ != "PING")
        if type_ == "PING":
            self.port_input.clear()

    def _add_target(self):
        target = self.target_input.text().strip()
        type_ = self.type_combo.currentText()
        port = self.port_input.text().strip()
        

        if not target:
            QMessageBox.critical(self, "错误", "请输入目标地址")
            return

        port_num = 0
        if type_ != "PING":
            if not port.isdigit() or not (0 < int(port) <= 65535):
                QMessageBox.critical(self, "错误", "端口必须为整数且范围在1-65535之间")
                return
            port_num = int(port)

        target_info = {
            "target": target,
            "type": type_,
            "port": port_num,
        }
        logging.debug(f"add_target_info:{target_info}")
        if self.model.add_target(target_info):
            self.target_input.clear()
            if type_ != "PING":
                self.port_input.clear()

    def start_detection(self):
        if self.active_tasks:
            return
        logging.debug(f"self.model._data:{self.model._data}")

        if self.model._data == []:
            QMessageBox.critical(self, "错误", "请先添加目标")
            return

        for row in range(self.model.rowCount()):
            target = self.model._data[row][0]
            type_ = self.model._data[row][1]
            port = 0 if self.model._data[row][2] == "-" else int(self.model._data[row][2])
            intreval = self.interval_input.text().strip()
            logging.debug(f'intreval:{intreval if intreval != "" else 1}')
            key = (target, type_, port,intreval)

            task = DetectionTask({
                "target": target,
                "type": type_,
                "port": port,
                "interval": intreval if intreval != "" else 1
            })
            task.signals.result_ready.connect(self.handle_result)
            self.active_tasks[key] = task
            self.thread_pool.start(task)

        self.start_btn.setEnabled(False)
        self.stop_btn.setEnabled(True)

    def stop_detection(self):
        for task in self.active_tasks.values():
            task.stop()
        self.active_tasks.clear()
        self.start_btn.setEnabled(True)
        self.stop_btn.setEnabled(False)

    def handle_result(self, key, success, delay):
        self.model.update_result(key, success, delay)
        self.csv_writer.add_record({
            "timestamp": datetime.now(),
            "target": key[0],
            "type": key[1],
            "port": key[2],
            "status": "SUCCESS" if success else "FAILURE",
            "delay": delay,
        })
   
    def delete_selected(self):
        selected = self.table_view.selectionModel().selectedRows()
        if not selected:
            QMessageBox.warning(self, "警告", "请先选择要删除的目标")
            return

        # 获取所有选中的行号
        rows = {index.row() for index in selected}
        
        # 停止相关检测任务
        for row in rows:
            target = self.model._data[row][0]
            type_ = self.model._data[row][1]
            port = 0 if self.model._data[row][2] == "-" else int(self.model._data[row][2])
           
            
            task_key = (target, type_, port)
            if task_key in self.active_tasks:
                self.active_tasks[task_key].stop()
                del self.active_tasks[task_key]

        # 从模型中删除数据
        self.model.remove_rows(rows)
        
        # 刷新表格视图
        self.table_view.viewport().update()
    def closeEvent(self, event):
        logging.info("关闭窗口")
        self.stop_detection()
        self.csv_writer.flush()
        event.accept()
    
    def reset_column_width(self):
        """重置列宽到默认值"""
        header = self.table_view.horizontalHeader()
        for col in range(header.count()):
            if header.sectionResizeMode(col) != QHeaderView.Stretch:
                header.resizeSection(col, self._get_default_width(col))
                
    def _get_default_width(self, column):
        """获取各列默认宽度"""
        defaults = {
            0: 250,  # 目标地址
            1: 80,   # 类型
            2: 70,   # 端口
            3: 100,  # 状态
            4: 120,  # 延迟
            5: 80,   # 成功
            6: 80,   # 总数
            7: 200   # 最后失败
        }
        return defaults.get(column, 100)


class CSVWriter:
    def __init__(self):
        self._buffer = []
        self._timer = QTimer()
        self._timer.timeout.connect(self.flush)
        self._timer.start(5000)

    def add_record(self, record):
        self._buffer.append(record)

    def flush(self):
        if not self._buffer:
            logging.info("No records to flush")
            return
        try:
            filename = datetime.now().strftime("network_log_%Y%m%d.csv")
            with open(filename, "a", newline="", encoding="utf-8") as f:
                writer = csv.DictWriter(f, fieldnames=[
                    "timestamp", "target", "type", "port", "status", "delay"
                ])
                if f.tell() == 0:
                    writer.writeheader()
                writer.writerows(self._buffer)
            self._buffer.clear()
        except Exception as e:
            logging.error(f"写入CSV文件失败: {str(e)}")
        
    
    def __del__(self):
        self.flush()

if __name__ == "__main__":
    app = QApplication(sys.argv)
    app.setStyle(QStyleFactory.create("Fusion"))
    window = MainWindow()
    window.show()
    sys.exit(app.exec())
